import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier  # 仅保留模型，移除可视化依赖
from sklearn.metrics import accuracy_score, classification_report
from joblib import dump

class ClassificationModel(object):
    def __init__(self):
        self.df = pd.read_csv('fina_indicator.csv')

    def get_conditions(self):
        df = self.df
        df['max_ratio'] = df['max_closes'] / df['the_closes']
        df['min_ratio'] = df['min_closes'] / df['the_closes']
        high_return_threshold = df['max_ratio'].quantile(0.4)
        high_risk_threshold = df['min_ratio'].quantile(0.4)

        conditions = [
            (df['max_ratio'] >= high_return_threshold) & (df['min_ratio'] >= high_risk_threshold),
            (df['max_ratio'] < high_return_threshold) & (df['min_ratio'] >= high_risk_threshold),
            (df['max_ratio'] >= high_return_threshold) & (df['min_ratio'] < high_risk_threshold),
            (df['max_ratio'] < high_return_threshold) & (df['min_ratio'] < high_risk_threshold)
        ]
        labels = ['高收益高风险', '低收益高风险', '高收益低风险', '低收益低风险']
        df['category'] = np.select(conditions, labels, default='未知')

        label_encoder = LabelEncoder()
        df['category_encoded'] = label_encoder.fit_transform(df['category'])

        scaler = StandardScaler()
        features = df[['eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin', 'fcff', 'fcfe',
                    'tangible_asset', 'bps', 'grossprofit_margin', 'npta']]
        X = scaler.fit_transform(features)
        y = df['category_encoded']
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=24)
        dump(scaler, 'scaler.joblib')
        dump(label_encoder,'label_encoder.joblib')
        feature_names = features.columns.tolist()
        return X_train, X_test, y_train, y_test, label_encoder, feature_names

    def knn_utils(self, X_train, X_test, y_train, y_test, label_encoder):
        knn = KNeighborsClassifier(n_neighbors=5)
        knn.fit(X_train, y_train)
        y_pred = knn.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        report = classification_report(y_test, y_pred, target_names=label_encoder.classes_)
        print(f"KNN 模型准确率: {accuracy:.4f}")
        print("KNN 分类报告:")
        print(report)
        dump(knn, 'knn.joblib')

    def svc_utils(self, X_train, X_test, y_train, y_test, label_encoder):
        svc = SVC(kernel='linear', random_state=24)
        svc.fit(X_train, y_train)
        y_pred = svc.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        report = classification_report(y_test, y_pred, target_names=label_encoder.classes_)
        print(f"SVC 模型准确率: {accuracy:.4f}")
        print("SVC 分类报告:")
        print(report)
        dump(svc, 'svc.joblib')

    def decision_tree_utils(self, X_train, X_test, y_train, y_test, label_encoder, feature_names):
        dt = DecisionTreeClassifier(
            max_depth=6,
            min_samples_split=10,
            min_samples_leaf=5,
            random_state=24,
            criterion='gini'
        )
        dt.fit(X_train, y_train)
        y_pred = dt.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        report = classification_report(y_test, y_pred, target_names=label_encoder.classes_)
        print(f"\n决策树 模型准确率: {accuracy:.4f}")
        print("决策树 分类报告:")
        print(report)
        dump(dt, 'decision_tree.joblib')
        print("决策树模型已保存为 decision_tree.joblib")

        feature_importance = pd.DataFrame({
            '特征名称': feature_names,
            '重要性': dt.feature_importances_
        }).sort_values('重要性', ascending=False)
        print("\n决策树 特征重要性排名:")
        print(feature_importance)

if __name__ == '__main__':
    cu = ClassificationModel()
    X_train, X_test, y_train, y_test, label_encoder, feature_names = cu.get_conditions()
    cu.decision_tree_utils(X_train, X_test, y_train, y_test, label_encoder, feature_names)
    # 若需运行其他分类器，取消注释即可
    cu.knn_utils(X_train, X_test, y_train, y_test, label_encoder)
    # cu.svc_utils(X_train, X_test, y_train, y_test, label_encoder)
